import os
import argparse
import pickle as pk
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

from stgcn import STGCN
from utils_recent import generate_dataset, load_metr_la_data, get_normalized_adj


# use_gpu = False
num_timesteps_input = 12
num_timesteps_output = 3


epochs = 200
batch_size = 32


def train_epoch(training_input, training_target, batch_size):
    """
    Trains one epoch with the given data.
    :param training_input: Training inputs of shape (num_samples, num_nodes,
    num_timesteps_train, num_features).
    :param training_target: Training targets of shape (num_samples, num_nodes,
    num_timesteps_predict).
    :param batch_size: Batch size to use during training.
    :return: Average loss for this epoch.
    """
    permutation = torch.randperm(training_input.shape[0])

    epoch_training_losses = []
    for i in range(0, training_input.shape[0], batch_size):
        net.train()
        optimizer.zero_grad()

        indices = permutation[i:i + batch_size]
        X_batch, y_batch = training_input[indices], training_target[indices]
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)
        out = net(A_wave, X_batch)
        loss = loss_criterion(out, y_batch)
        loss.backward()
        optimizer.step()
        epoch_training_losses.append(loss.detach().cpu().numpy())
        #print(loss)
    return sum(epoch_training_losses)/len(epoch_training_losses)


if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    A, X, means, stds = load_metr_la_data()

    split_line1 = int(X.shape[2] * 0.6)
    split_line2 = int(X.shape[2] * 0.8)

    train_original_data = X[:, :, :split_line1]
    val_original_data = X[:, :, split_line1:]
    test_original_data = X[:, :, split_line2:]


    training_input, training_target = generate_dataset(train_original_data,
                                                       num_timesteps_input=num_timesteps_input,
                                                       num_timesteps_output=num_timesteps_output)

    val_input, val_target = generate_dataset(train_original_data,
                                             num_timesteps_input=num_timesteps_input,
                                             num_timesteps_output=num_timesteps_output)

    test_input, test_target = generate_dataset(test_original_data,
                                               num_timesteps_input=num_timesteps_input,
                                               num_timesteps_output=num_timesteps_output)

    A_wave = get_normalized_adj(A)
    A_wave = torch.from_numpy(A_wave)

    A_wave = A_wave.to(device)

    net = STGCN(A_wave.shape[0],
                training_input.shape[3],
                num_timesteps_input,
                num_timesteps_output).to(device)

    optimizer = torch.optim.Adam(net.parameters(), lr=1e-3,weight_decay=0.001)
    loss_criterion = nn.MSELoss()

    training_losses = []
    validation_losses = []
    validation_maes = []
    validation_mape = []

    for epoch in range(epochs):
        loss = train_epoch(training_input, training_target,
                           batch_size=batch_size)
        training_losses.append(loss)


        # Run validation
        with torch.no_grad():
            net.eval()
            val_input = val_input.to(device)
            val_target = val_target.to(device)

            out = net(A_wave, val_input)
            val_loss = loss_criterion(out, val_target).to(device)
            validation_losses.append(np.asscalar(val_loss.detach().cpu().numpy()))

            out_unnormalized = out.detach().cpu().numpy() * stds[0] + means[0]
            target_unnormalized = val_target.detach().cpu().numpy() * stds[0] + means[0]

            mae = np.mean(np.absolute(out_unnormalized - target_unnormalized))
            validation_maes.append(mae)
            mape = np.mean(np.abs((out_unnormalized - target_unnormalized) / (target_unnormalized + 10))) * 100
            validation_mape.append(mape)
            out = None
            val_input = val_input.to(device)
            val_target = val_target.to(device)

        print("Training loss:{} ".format(training_losses))
        print("Validation loss: {}".format(validation_losses))
        print("Validation MAE: {}".format(validation_maes))
        print("Validation MAPE: {}".format(validation_mape))
